//
// Created by yk120 on 2024/3/1.
//

#include <nanofaiss/IndexIVF.h>
#include <nanofaiss/impl/FaissException.h>
#include <nanofaiss/impl/FaissAssert.h>

#include <memory>

namespace faiss {


Level1Quantizer::Level1Quantizer(Index* quantizer, size_t nlist)
        : quantizer(quantizer), nlist(nlist) {
    cp.niter = 10; // set low iterations
}

Level1Quantizer::Level1Quantizer() = default;


Level1Quantizer::~Level1Quantizer() {
    if (own_fields)
        delete quantizer;
}

void Level1Quantizer::train_q1(size_t n, const float* x, MetricType metric) {
    size_t d = quantizer->d;

    /// use quantizer as the cluster centroid, and every point
    /// in the quantizer is seen an inverted list or cluster
    if (quantizer->is_trained && quantizer->ntotal == nlist) {

    } else if (quantizer_train_alone == 1) {

    } else if (quantizer_train_alone == 0) { /// use quantizer as the kmeans index, save the cluster in ther quantizer
        Clustering cluster(d, nlist, cp);
        quantizer->reset(); // reset the index of quantizer

        if (clustering_index) {

        } else {
            cluster.train(n, x, *quantizer);
        }
        quantizer->is_trained = true;
    } else if (quantizer_train_alone == 2) {
        //TODO
    }
}

IndexIVF::IndexIVF(Index* quantizer,
                   size_t d,
                   size_t nlist,
                   size_t code_size,
                   MetricType metric)
                    : Index(d, metric), IndexIVFInterface(quantizer, nlist),
                      invlists(new ArrayInvertedLists(nlist, code_size)),
                      own_invlists(true), code_size(code_size) {
    FAISS_THROW_IF_NOT(d == quantizer->d);
    is_trained = quantizer->is_trained && (quantizer->ntotal == nlist);

    if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
        // TODO
    }

}

idx_t IndexIVF::train_encoder_num_vectors() const {
    return 0;
}

void IndexIVF::train_encoder(idx_t n, const float* x, const idx_t* assign) {
    // pass
}


void IndexIVF::train(idx_t n, const float* x) {
    // train the inverted list
    train_q1(n, x, metric_type);


    idx_t max_nt = train_encoder_num_vectors();
    if (max_nt <= 0) {
        max_nt = static_cast<size_t>(1 << 35);
    }

    TransformedVectors tv(x, x); // TODO: May subsample

    if (by_residual) {
        //TODO for residual
    } else {
        train_encoder(n, tv.x, nullptr);
    }

    is_trained = true;
}

void IndexIVF::add(faiss::idx_t n, const float* x) {
    add_with_ids(n, x, nullptr);
}

void IndexIVF::add_with_ids(idx_t n, const float* x, const idx_t* xids) {
    std::unique_ptr<idx_t[]> course_idx(new idx_t[n]);
    quantizer->assign(n, x, course_idx.get());
    add_core(n, x, xids, course_idx.get());
}




}